﻿// Copyright (c) Microsoft Corporation.  All rights reserved.
// This source code is made available under the terms of the Microsoft Public License (MS-PL)

using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using System.Linq.Expressions;
using System.Reflection;

namespace CoolCode.Linq.Expressions {
	public static class ExpressionVisitorExtensions {
		public static Expression Visit<T>(
			this Expression exp,
			Func<T, Expression> visitor) where T : Expression {
			return ExpressionVisitor<T>.Visit(exp, visitor);
		}

		public static TExp Visit<T, TExp>(
			this TExp exp,
			Func<T, Expression> visitor)
			where T : Expression
			where TExp : Expression {
			return (TExp)ExpressionVisitor<T>.Visit(exp, visitor);
		}

		public static Expression<TDelegate> Visit<T, TDelegate>(
			this Expression<TDelegate> exp,
			Func<T, Expression> visitor) where T : Expression {
			return ExpressionVisitor<T>.Visit<TDelegate>(exp, visitor);
		}

		public static IQueryable<TSource> Visit<T, TSource>(
			this IQueryable<TSource> source,
			Func<T, Expression> visitor) where T : Expression {
			return source.Provider.CreateQuery<TSource>(ExpressionVisitor<T>.Visit(source.Expression, visitor));
		}
	}

	/// <summary>
	/// This class visits every Parameter expression in an expression tree and calls a delegate
	/// to optionally replace the parameter.  This is useful where two expression trees need to
	/// be merged (and they don't share the same ParameterExpressions).
	/// </summary>
	public class ExpressionVisitor<T> : ExpressionVisitor where T : Expression {
		Func<T, Expression> visitor;

		public ExpressionVisitor(Func<T, Expression> visitor) {
			this.visitor = visitor;
		}

		public static Expression Visit(
			Expression exp,
			Func<T, Expression> visitor) {
			return new ExpressionVisitor<T>(visitor).Visit(exp);
		}

		public static Expression<TDelegate> Visit<TDelegate>(
			Expression<TDelegate> exp,
			Func<T, Expression> visitor) {
			return (Expression<TDelegate>)new ExpressionVisitor<T>(visitor).Visit(exp);
		}

		protected override Expression Visit(Expression exp) {
			if (exp is T && visitor != null)
				exp = visitor((T)exp);

			return base.Visit(exp);
		}
	}

	public abstract class ExpressionVisitor {
		protected ExpressionVisitor() {
		}

		protected virtual Expression Visit(Expression exp) {
			if (exp == null)
				return exp;
			switch (exp.NodeType) {
				case ExpressionType.Negate:
				case ExpressionType.NegateChecked:
				case ExpressionType.Not:
				case ExpressionType.Convert:
				case ExpressionType.ConvertChecked:
				case ExpressionType.ArrayLength:
				case ExpressionType.Quote:
				case ExpressionType.TypeAs:
				case ExpressionType.UnaryPlus:
					return this.VisitUnary((UnaryExpression)exp);
				case ExpressionType.Add:
				case ExpressionType.AddChecked:
				case ExpressionType.Subtract:
				case ExpressionType.SubtractChecked:
				case ExpressionType.Multiply:
				case ExpressionType.MultiplyChecked:
				case ExpressionType.Divide:
				case ExpressionType.Modulo:
				case ExpressionType.And:
				case ExpressionType.AndAlso:
				case ExpressionType.Or:
				case ExpressionType.OrElse:
				case ExpressionType.LessThan:
				case ExpressionType.LessThanOrEqual:
				case ExpressionType.GreaterThan:
				case ExpressionType.GreaterThanOrEqual:
				case ExpressionType.Equal:
				case ExpressionType.NotEqual:
				case ExpressionType.Coalesce:
				case ExpressionType.ArrayIndex:
				case ExpressionType.RightShift:
				case ExpressionType.LeftShift:
				case ExpressionType.ExclusiveOr:
				case ExpressionType.Power:
					return this.VisitBinary((BinaryExpression)exp);
				case ExpressionType.TypeIs:
					return this.VisitTypeIs((TypeBinaryExpression)exp);
				case ExpressionType.Conditional:
					return this.VisitConditional((ConditionalExpression)exp);
				case ExpressionType.Constant:
					return this.VisitConstant((ConstantExpression)exp);
				case ExpressionType.Parameter:
					return this.VisitParameter((ParameterExpression)exp);
				case ExpressionType.MemberAccess:
					return this.VisitMemberAccess((MemberExpression)exp);
				case ExpressionType.Call:
					return this.VisitMethodCall((MethodCallExpression)exp);
				case ExpressionType.Lambda:
					return this.VisitLambda((LambdaExpression)exp);
				case ExpressionType.New:
					return this.VisitNew((NewExpression)exp);
				case ExpressionType.NewArrayInit:
				case ExpressionType.NewArrayBounds:
					return this.VisitNewArray((NewArrayExpression)exp);
				case ExpressionType.Invoke:
					return this.VisitInvocation((InvocationExpression)exp);
				case ExpressionType.MemberInit:
					return this.VisitMemberInit((MemberInitExpression)exp);
				case ExpressionType.ListInit:
					return this.VisitListInit((ListInitExpression)exp);
				default:
					return this.VisitUnknown(exp);
			}
		}

		protected virtual Expression VisitUnknown(Expression expression) {
			throw new Exception(string.Format("Unhandled expression type: '{0}'", expression.NodeType));
		}

		protected virtual MemberBinding VisitBinding(MemberBinding binding) {
			switch (binding.BindingType) {
				case MemberBindingType.Assignment:
					return this.VisitMemberAssignment((MemberAssignment)binding);
				case MemberBindingType.MemberBinding:
					return this.VisitMemberMemberBinding((MemberMemberBinding)binding);
				case MemberBindingType.ListBinding:
					return this.VisitMemberListBinding((MemberListBinding)binding);
				default:
					throw new Exception(string.Format("Unhandled binding type '{0}'", binding.BindingType));
			}
		}

		protected virtual ElementInit VisitElementInitializer(ElementInit initializer) {
			ReadOnlyCollection<Expression> arguments = this.VisitExpressionList(initializer.Arguments);
			if (arguments != initializer.Arguments) {
				return Expression.ElementInit(initializer.AddMethod, arguments);
			}
			return initializer;
		}

		protected virtual Expression VisitUnary(UnaryExpression u) {
			Expression operand = this.Visit(u.Operand);
			return this.UpdateUnary(u, operand, u.Type, u.Method);
		}

		protected UnaryExpression UpdateUnary(UnaryExpression u, Expression operand, Type resultType, MethodInfo method) {
			if (u.Operand != operand || u.Type != resultType || u.Method != method) {
				return Expression.MakeUnary(u.NodeType, operand, resultType, method);
			}
			return u;
		}

		protected virtual Expression VisitBinary(BinaryExpression b) {
			Expression left = this.Visit(b.Left);
			Expression right = this.Visit(b.Right);
			Expression conversion = this.Visit(b.Conversion);
			return this.UpdateBinary(b, left, right, conversion, b.IsLiftedToNull, b.Method);
		}

		protected BinaryExpression UpdateBinary(BinaryExpression b, Expression left, Expression right, Expression conversion, bool isLiftedToNull, MethodInfo method) {
			if (left != b.Left || right != b.Right || conversion != b.Conversion || method != b.Method || isLiftedToNull != b.IsLiftedToNull) {
				if (b.NodeType == ExpressionType.Coalesce && b.Conversion != null) {
					return Expression.Coalesce(left, right, conversion as LambdaExpression);
				}
				else {
					return Expression.MakeBinary(b.NodeType, left, right, isLiftedToNull, method);
				}
			}
			return b;
		}

		protected virtual Expression VisitTypeIs(TypeBinaryExpression b) {
			Expression expr = this.Visit(b.Expression);
			return this.UpdateTypeIs(b, expr, b.TypeOperand);
		}

		protected TypeBinaryExpression UpdateTypeIs(TypeBinaryExpression b, Expression expression, Type typeOperand) {
			if (expression != b.Expression || typeOperand != b.TypeOperand) {
				return Expression.TypeIs(expression, typeOperand);
			}
			return b;
		}

		protected virtual Expression VisitConstant(ConstantExpression c) {
			return c;
		}

		protected virtual Expression VisitConditional(ConditionalExpression c) {
			Expression test = this.Visit(c.Test);
			Expression ifTrue = this.Visit(c.IfTrue);
			Expression ifFalse = this.Visit(c.IfFalse);
			return this.UpdateConditional(c, test, ifTrue, ifFalse);
		}

		protected ConditionalExpression UpdateConditional(ConditionalExpression c, Expression test, Expression ifTrue, Expression ifFalse) {
			if (test != c.Test || ifTrue != c.IfTrue || ifFalse != c.IfFalse) {
				return Expression.Condition(test, ifTrue, ifFalse);
			}
			return c;
		}

		protected virtual Expression VisitParameter(ParameterExpression p) {
			return p;
		}

		protected virtual Expression VisitMemberAccess(MemberExpression m) {
			Expression exp = this.Visit(m.Expression);
			return this.UpdateMemberAccess(m, exp, m.Member);
		}

		protected MemberExpression UpdateMemberAccess(MemberExpression m, Expression expression, MemberInfo member) {
			if (expression != m.Expression || member != m.Member) {
				return Expression.MakeMemberAccess(expression, member);
			}
			return m;
		}

		protected virtual Expression VisitMethodCall(MethodCallExpression m) {
			Expression obj = this.Visit(m.Object);
			IEnumerable<Expression> args = this.VisitExpressionList(m.Arguments);
			return this.UpdateMethodCall(m, obj, m.Method, args);
		}

		protected MethodCallExpression UpdateMethodCall(MethodCallExpression m, Expression obj, MethodInfo method, IEnumerable<Expression> args) {
			if (obj != m.Object || method != m.Method || args != m.Arguments) {
				return Expression.Call(obj, method, args);
			}
			return m;
		}

		protected virtual ReadOnlyCollection<Expression> VisitExpressionList(ReadOnlyCollection<Expression> original) {
			if (original != null) {
				List<Expression> list = null;
				for (int i = 0, n = original.Count; i < n; i++) {
					Expression p = this.Visit(original[i]);
					if (list != null) {
						list.Add(p);
					}
					else if (p != original[i]) {
						list = new List<Expression>(n);
						for (int j = 0; j < i; j++) {
							list.Add(original[j]);
						}
						list.Add(p);
					}
				}
				if (list != null) {
					return list.AsReadOnly();
				}
			}
			return original;
		}

		protected virtual MemberAssignment VisitMemberAssignment(MemberAssignment assignment) {
			Expression e = this.Visit(assignment.Expression);
			return this.UpdateMemberAssignment(assignment, assignment.Member, e);
		}

		protected MemberAssignment UpdateMemberAssignment(MemberAssignment assignment, MemberInfo member, Expression expression) {
			if (expression != assignment.Expression || member != assignment.Member) {
				return Expression.Bind(member, expression);
			}
			return assignment;
		}

		protected virtual MemberMemberBinding VisitMemberMemberBinding(MemberMemberBinding binding) {
			IEnumerable<MemberBinding> bindings = this.VisitBindingList(binding.Bindings);
			return this.UpdateMemberMemberBinding(binding, binding.Member, bindings);
		}

		protected MemberMemberBinding UpdateMemberMemberBinding(MemberMemberBinding binding, MemberInfo member, IEnumerable<MemberBinding> bindings) {
			if (bindings != binding.Bindings || member != binding.Member) {
				return Expression.MemberBind(member, bindings);
			}
			return binding;
		}

		protected virtual MemberListBinding VisitMemberListBinding(MemberListBinding binding) {
			IEnumerable<ElementInit> initializers = this.VisitElementInitializerList(binding.Initializers);
			return this.UpdateMemberListBinding(binding, binding.Member, initializers);
		}

		protected MemberListBinding UpdateMemberListBinding(MemberListBinding binding, MemberInfo member, IEnumerable<ElementInit> initializers) {
			if (initializers != binding.Initializers || member != binding.Member) {
				return Expression.ListBind(member, initializers);
			}
			return binding;
		}

		protected virtual IEnumerable<MemberBinding> VisitBindingList(ReadOnlyCollection<MemberBinding> original) {
			List<MemberBinding> list = null;
			for (int i = 0, n = original.Count; i < n; i++) {
				MemberBinding b = this.VisitBinding(original[i]);
				if (list != null) {
					list.Add(b);
				}
				else if (b != original[i]) {
					list = new List<MemberBinding>(n);
					for (int j = 0; j < i; j++) {
						list.Add(original[j]);
					}
					list.Add(b);
				}
			}
			if (list != null)
				return list;
			return original;
		}

		protected virtual IEnumerable<ElementInit> VisitElementInitializerList(ReadOnlyCollection<ElementInit> original) {
			List<ElementInit> list = null;
			for (int i = 0, n = original.Count; i < n; i++) {
				ElementInit init = this.VisitElementInitializer(original[i]);
				if (list != null) {
					list.Add(init);
				}
				else if (init != original[i]) {
					list = new List<ElementInit>(n);
					for (int j = 0; j < i; j++) {
						list.Add(original[j]);
					}
					list.Add(init);
				}
			}
			if (list != null)
				return list;
			return original;
		}

		protected virtual Expression VisitLambda(LambdaExpression lambda) {
			Expression body = this.Visit(lambda.Body);
			return this.UpdateLambda(lambda, lambda.Type, body, lambda.Parameters);
		}

		protected LambdaExpression UpdateLambda(LambdaExpression lambda, Type delegateType, Expression body, IEnumerable<ParameterExpression> parameters) {
			if (body != lambda.Body || parameters != lambda.Parameters || delegateType != lambda.Type) {
				return Expression.Lambda(delegateType, body, parameters);
			}
			return lambda;
		}

		protected virtual NewExpression VisitNew(NewExpression nex) {
			IEnumerable<Expression> args = this.VisitExpressionList(nex.Arguments);
			return this.UpdateNew(nex, nex.Constructor, args, nex.Members);
		}

		protected NewExpression UpdateNew(NewExpression nex, ConstructorInfo constructor, IEnumerable<Expression> args, IEnumerable<MemberInfo> members) {
			if (args != nex.Arguments || constructor != nex.Constructor || members != nex.Members) {
				if (nex.Members != null) {
					return Expression.New(constructor, args, members);
				}
				else {
					return Expression.New(constructor, args);
				}
			}
			return nex;
		}

		protected virtual Expression VisitMemberInit(MemberInitExpression init) {
			NewExpression n = this.VisitNew(init.NewExpression);
			IEnumerable<MemberBinding> bindings = this.VisitBindingList(init.Bindings);
			return this.UpdateMemberInit(init, n, bindings);
		}

		protected MemberInitExpression UpdateMemberInit(MemberInitExpression init, NewExpression nex, IEnumerable<MemberBinding> bindings) {
			if (nex != init.NewExpression || bindings != init.Bindings) {
				return Expression.MemberInit(nex, bindings);
			}
			return init;
		}

		protected virtual Expression VisitListInit(ListInitExpression init) {
			NewExpression n = this.VisitNew(init.NewExpression);
			IEnumerable<ElementInit> initializers = this.VisitElementInitializerList(init.Initializers);
			return this.UpdateListInit(init, n, initializers);
		}

		protected ListInitExpression UpdateListInit(ListInitExpression init, NewExpression nex, IEnumerable<ElementInit> initializers) {
			if (nex != init.NewExpression || initializers != init.Initializers) {
				return Expression.ListInit(nex, initializers);
			}
			return init;
		}

		protected virtual Expression VisitNewArray(NewArrayExpression na) {
			IEnumerable<Expression> exprs = this.VisitExpressionList(na.Expressions);
			return this.UpdateNewArray(na, na.Type, exprs);
		}

		protected NewArrayExpression UpdateNewArray(NewArrayExpression na, Type arrayType, IEnumerable<Expression> expressions) {
			if (expressions != na.Expressions || na.Type != arrayType) {
				if (na.NodeType == ExpressionType.NewArrayInit) {
					return Expression.NewArrayInit(arrayType.GetElementType(), expressions);
				}
				else {
					return Expression.NewArrayBounds(arrayType.GetElementType(), expressions);
				}
			}
			return na;
		}

		protected virtual Expression VisitInvocation(InvocationExpression iv) {
			IEnumerable<Expression> args = this.VisitExpressionList(iv.Arguments);
			Expression expr = this.Visit(iv.Expression);
			return this.UpdateInvocation(iv, expr, args);
		}

		protected InvocationExpression UpdateInvocation(InvocationExpression iv, Expression expression, IEnumerable<Expression> args) {
			if (args != iv.Arguments || expression != iv.Expression) {
				return Expression.Invoke(expression, args);
			}
			return iv;
		}
	}
}